package com.cicc.wm.lock.jdbc;

import com.cicc.wm.lock.DistributedLock;
import com.cicc.wm.lock.LockRepository;
import com.cicc.wm.lock.WheelTimerService;
import io.netty.util.Timeout;
import io.netty.util.TimerTask;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.dao.CannotAcquireLockException;
import org.springframework.dao.DataAccessResourceFailureException;
import org.springframework.dao.TransientDataAccessException;
import org.springframework.transaction.TransactionSystemException;
import org.springframework.transaction.TransactionTimedOutException;

import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * JDBC分布式锁实现
 * 支持可重入、自动续期、超时自动释放等特性
 */
@Data
@Slf4j
public class JdbcDistributedLock implements DistributedLock {
    /**
     * 锁续期任务映射表
     */
    private static final Map<String, ExpirationEntry> EXPIRATION_RENEWAL_MAP = new ConcurrentHashMap<>();

    /**
     * 获取锁重试时间间隔(毫秒)
     */
    private static final int DEFAULT_IDLE = 100;

    /**
     * 默认锁超时时间
     */
    private final Duration defaultTtl = Duration.ofSeconds(30);

    /**
     * 锁仓库
     */
    private final LockRepository mutex;

    /**
     * 锁键
     */
    private final String lockKey;

    /**
     * 锁租期时间(毫秒)
     */
    private long leaseTime = -1;

    public JdbcDistributedLock(LockRepository mutex, String lockKey) {
        this.mutex = mutex;
        this.lockKey = lockKey;
    }

    @Override
    public void lock() {
        lock(-1, null);
    }

    @Override
    public void lock(long leaseTime, TimeUnit unit) {
        while (true) {
            try {
                while (!doLock(leaseTime, unit)) {
                    TimeUnit.MILLISECONDS.sleep(DEFAULT_IDLE);
                }
                break;
            } catch (InterruptedException e) {
                /*
                 * This method must be uninterruptible so catch and ignore
                 * interrupts and only break out of the while loop when
                 * we get the lock.
                 */
            } catch (TransientDataAccessException | TransactionTimedOutException | TransactionSystemException e) {
                // try again
            } catch (Exception e) {
                rethrowAsLockException(e);
            }
        }
    }

    @Override
    public boolean tryLock() {
        try {
            return tryLock(0, TimeUnit.MICROSECONDS);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            return false;
        }
    }

    @Override
    public boolean tryLock(long waitTime, TimeUnit timeUnit) throws InterruptedException {
        return this.tryLock(waitTime, -1, timeUnit);
    }

    @Override
    public boolean tryLock(long waitTime, long leaseTime, TimeUnit unit) throws InterruptedException {
        long now = System.currentTimeMillis();
        long expire = now + TimeUnit.MILLISECONDS.convert(waitTime, unit);
        boolean acquired;
        while (true) {
            try {
                do {
                    acquired = doLock(leaseTime, unit);
                    if (!acquired) {
                        // 释放锁
                        TimeUnit.MILLISECONDS.sleep(DEFAULT_IDLE);
                    }
                } while (!acquired && System.currentTimeMillis() < expire);
                return acquired;
            } catch (TransientDataAccessException | TransactionTimedOutException | TransactionSystemException e) {
                // try again
            } catch (Exception e) {
                rethrowAsLockException(e);
            }
        }
    }

    private boolean doLock(long leaseTime, TimeUnit unit) {
        long lockLeaseTime = leaseTime > 0 ? unit.toMillis(leaseTime) : this.defaultTtl.toMillis();
        boolean acquired = this.mutex.acquire(Thread.currentThread().threadId(), this.lockKey, lockLeaseTime);

        if (acquired) {
            this.leaseTime = leaseTime > 0 ? unit.toMillis(leaseTime) : leaseTime;
            if (this.leaseTime > 0) {
                scheduleExpirationRelease(this.lockKey, leaseTime, unit);
            } else {
                scheduleExpirationRenewal(Thread.currentThread().threadId());
            }
        }

        return acquired;
    }

    @Override
    public void unlock() {
        try {
            long threadId = Thread.currentThread().threadId();
            if (this.mutex.isAcquired(threadId, this.lockKey)) {
                if (this.mutex.delete(threadId, this.lockKey)) {
                    this.cancelExpirationRenewal();
                } else {
                    log.warn("Failed to release lock for thread {} on key {}", threadId, this.lockKey);
                }
            }
        } catch (Exception e) {
            throw new DataAccessResourceFailureException(
                    String.format("Failed to release mutex at %s", this.lockKey), e);
        }
    }

    private void rethrowAsLockException(Exception e) {
        throw new CannotAcquireLockException("Failed to lock mutex at " + this.lockKey, e);
    }

    /**
     * 安排锁自动释放任务
     *
     * @param lockKey 锁key
     * @param timeout 释放锁时间
     * @param unit    时间单位
     */
    private void scheduleExpirationRelease(String lockKey, long timeout, TimeUnit unit) {
        long threadId = Thread.currentThread().threadId();
        WheelTimerService.newTimeout(t -> {
            try {
                if (this.mutex.delete(threadId, lockKey)) {
                    log.debug("Lock {} automatically released for thread {}", lockKey, threadId);
                }
            } catch (Exception e) {
                log.error("Failed to release lock at {} for thread {}", lockKey, threadId, e);
            }
        }, timeout, unit);
    }

    protected void scheduleExpirationRenewal(long threadId) {
        ExpirationEntry entry = new ExpirationEntry();
        ExpirationEntry oldEntry = EXPIRATION_RENEWAL_MAP.putIfAbsent(this.getLockKey(), entry);

        if (oldEntry != null) {
            oldEntry.setReentryCount(oldEntry.getReentryCount() + 1);
            oldEntry.setThreadId(threadId);
        } else {
            entry.setThreadId(threadId);
            renewExpiration();
        }
    }

    /**
     * 锁自动续期任务
     */
    private void renewExpiration() {
        ExpirationEntry expirationEntry = EXPIRATION_RENEWAL_MAP.get(this.getLockKey());
        if (expirationEntry == null) {
            return;
        }

        Timeout task = WheelTimerService.newTimeout(new TimerTask() {
            @Override
            public void run(Timeout timeout) throws Exception {
                ExpirationEntry ent = EXPIRATION_RENEWAL_MAP.get(JdbcDistributedLock.this.getLockKey());
                if (ent == null) {
                    return;
                }

                if (renewLock(ent.getThreadId())) {
                    renewExpiration();
                } else {
                    log.warn("Failed to renew lock {} for thread {}",
                            JdbcDistributedLock.this.getLockKey(), ent.getThreadId());
                }
            }
        }, defaultTtl.toSeconds() / 3L, TimeUnit.SECONDS);

        expirationEntry.setTimeout(task);
    }

    protected void cancelExpirationRenewal() {
        ExpirationEntry expirationEntry = EXPIRATION_RENEWAL_MAP.get(this.getLockKey());
        if (expirationEntry == null) {
            return;
        }

        int reentryCount = expirationEntry.getReentryCount() - 1;
        if (reentryCount <= 0) {
            Timeout timeout = expirationEntry.getTimeout();
            if (timeout != null) {
                timeout.cancel();
                EXPIRATION_RENEWAL_MAP.remove(this.getLockKey());
            }
        }
    }

    /**
     * 锁续期实现
     *
     * @param threadId 持锁线程ID
     * @return 返回续期结果
     */
    private boolean renewLock(long threadId) {
        try {
            return this.mutex.renew(threadId, this.lockKey);
        } catch (Exception e) {
            log.error("Failed to renew lock at " + this.lockKey, e);
            return false;
        }
    }

    /**
     * 过期实体信息
     */
    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    private static class ExpirationEntry {
        /**
         * 持锁线程ID
         */
        private Long threadId;
        /**
         * 重入次数
         */
        private Integer reentryCount = 1;

        /**
         * 续期任务
         */
        private Timeout timeout;
    }
}
